/**
 * Copyright 2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifndef AICPU_KERNEL_BUILDER_H_
#define AICPU_KERNEL_BUILDER_H_

#include <nlohmann/json.hpp>

#include "aicpu_ops_kernel_info_store/op_struct.h"
#include "aicpu_ops_kernel_builder/aicpu_ops_kernel_builder.h"
#include "error_code/error_code.h"
#include "factory/factory.h"
#include "register/ops_kernel_builder_registry.h"
namespace aicpu {

class KernelBuilder {
 public:
  /**
   * constructor
   * @param void
   */
  KernelBuilder() = default;
  /**
   * Destructor
   */
  virtual ~KernelBuilder() = default;

  /**
   * init kernel builder
   * @return status whether this operation success
   */
  virtual ge::Status Initialize() { return ge::SUCCESS; }
  /**
   * Release related resources of the aicpu kernel builder
   * @return status whether this operation success
   */
  virtual ge::Status Finalize() { return ge::SUCCESS; }

  /**
   * Calc the running size of Operator,then GE will alloc the memsize from
   * runtime The size is consist of the part as follow: 1.StrFWKKernel; 2.Input
   * and output size; 3.NodeDef in tf; 4.FuncDef in tf.
   * @param node Node information, return task_memsize in node's attr
   * @return status whether this operation success
   */
  virtual ge::Status CalcOpRunningParam(const ge::Node &node) const = 0;

  /**
   * Calc the running size of Operator,then GE will alloc the memsize from
   * runtime The size is consist of the part as follow: 1.StrFWKKernel; 2.Input
   * and output size; 3.NodeDef in tf; 4.FuncDef in tf.
   * @param node Node information, return task_memsize in node's attr
   * @return status whether this operation success
   */
  virtual ge::Status GenerateTask(const ge::Node &node,
                                  const ge::RunContext &context,
                                  std::vector<domi::TaskDef> &tasks) = 0;
  /**
   * Generate the task
   * @param node Node information
   * @param task[out]
   * @param task_info[out]
   * @return status whether this operation success
   */
  virtual ge::Status GenSingleOpRunTask(const ge::NodePtr &node,
                                        STR_FWK_OP_KERNEL &task,
                                        std::string &task_info) {
    return ge::FAILED;
  };

  /**
   * Generate the task
   * @param count the memcopy times
   * @param task[out]
   * @param task_info[out]
   * @return status whether this operation success
   */
  virtual ge::Status GenMemCopyTask(uint64_t count, STR_FWK_OP_KERNEL &task,
                                    std::string &task_info) {
    return ge::FAILED;
  };

 protected:
  /**
   * Get the total size by datatype and ge shape
   * @param data_type Represent that how many size for each data
   * @param ge_shape Represent the dims information
   * @param total_size The total size
   * @return whether handle success
   */
  State GetTotalSizeByShapeAndType(const ge::DataType &data_type,
                                   const ge::GeShape &ge_shape,
                                   int64_t &total_size) const;
  /**
   * Get workspace info
   * @param op_desc_ptr Ge op description pointer
   * @param data_mem_base Data memory base addr
   * @param workspace_bytes_size Workspace bytes size
   * @return whether handle success
   */
  ge::Status GetWorkspaceInfo(const ge::OpDescPtr &op_desc_ptr,
                              uint8_t *data_mem_base,
                              uint64_t &workspace_bytes_size) const;

  /**
   * Calculate and set outputs size
   * @param op_desc_ptr Ge op description pointer
   * @return whether handle success
   */
  ge::Status SetOutPutsSize(std::shared_ptr<ge::OpDesc> &op_desc_ptr) const;

  /**
   * Set outputs size for unknow type 4 which output is result summary
   * @param op_desc_ptr Ge op description pointer
   * @return whether handle success
   */
  ge::Status SetOutSizeForSummary(std::shared_ptr<ge::OpDesc> &op_desc_ptr) const;

  /**
   * Calculate output size for unknow type 3 which output shape is range
   * @param data_type Represent that how many size for each data
   * @param shape_range outputs shape range for unknow type 3
   * @param total_size output total size
   * @return whether handle success
   */
  ge::Status GetOutSizeByShapeRange(
      const ge::DataType &data_type,
      const std::vector<std::pair<int64_t, int64_t>> &shape_range,
      int64_t &total_size) const;

  /**
   * Make extend info for op name: must be the first extend info (RUNTIME only
   * decipher the first extend info)
   * @param op_desc_ptr Ge op description pointer
   * @param tastExtInfo task extend info
   * @return whether handle success
   */
  ge::Status MakeExtInfoForOpName(const ge::OpDescPtr &op_desc_ptr,
                                  std::vector<char> &task_ext_info) const;
  /**
   * Make common task extend info
   * @param op_desc_ptr Ge op description pointer
   * @param tastExtInfo task extend info
   * @return whether handle success
   */
  ge::Status MakeBaseExtInfo(const ge::OpDescPtr &op_desc_ptr,
                             std::vector<char> &task_ext_info) const;

  virtual void GetInOutPutsDataType(const ge::OpDescPtr &op_desc_ptr,
                                    std::vector<uint32_t> &inputs_type,
                                    std::vector<uint32_t> &outputs_type) const {}
};

#define FACTORY_KERNEL_BUILDER Factory<KernelBuilder>

#define FACTORY_KERNEL_BUILDER_CLASS_KEY(CLASS, KEY) \
  FACTORY_KERNEL_BUILDER::Register<CLASS> __##CLASS(KEY);

}  // namespace aicpu

#endif  // AICPU_KERNEL_BUILDER_H_
